#ifndef GMM__H
#define GMM__H

#include <Eigen/Core>
#include <iostream>
#include <pcl/common/transforms.h>
#include <pcl/filters/crop_box.h>
#include <pcl/filters/passthrough.h>
#include <pcl/filters/voxel_grid.h>
#include <pcl/pcl_base.h>
#include <pcl/segmentation/extract_clusters.h>
#include <queue>
#include <random>
#include <ros/ros.h>

typedef pcl::PointXYZ POINT_TYPE;
typedef pcl::PointCloud<POINT_TYPE> PC_TYPE;
typedef PC_TYPE::Ptr PC_PTR_TYPE;

namespace gvm {

struct normal_random_variable {
  normal_random_variable(Eigen::MatrixXd const &covar)
      : normal_random_variable(Eigen::VectorXd::Zero(covar.rows()), covar) {}

  normal_random_variable(Eigen::VectorXd const &mean,
                         Eigen::MatrixXd const &covar)
      : mean(mean) {
    Eigen::SelfAdjointEigenSolver<Eigen::MatrixXd> eigenSolver(covar);
    transform = eigenSolver.eigenvectors() *
                eigenSolver.eigenvalues().cwiseSqrt().asDiagonal();
  }

  Eigen::VectorXd operator()() const {
    static std::mt19937 gen{std::random_device{}()};
    static std::normal_distribution<> dist;
    return mean + transform * Eigen::VectorXd{mean.size()}.unaryExpr(
                                  [&](double x) { return dist(gen); });
  }
  Eigen::VectorXd mean;
  Eigen::MatrixXd transform;
};

class GMM {
public:
  GMM() {}
  ~GMM() {}
  template <typename T_VEC_DOUBLE, typename T_VEC_EIGEN_VECTOR,
            typename T_VEC_EIGEN_MATRIX>
  GMM(T_VEC_DOUBLE &&pi_vector, T_VEC_EIGEN_VECTOR &&means,
      T_VEC_EIGEN_MATRIX &&var, int n_component,
      Eigen::Vector3d sensor_location)
      : pi_vector_(pi_vector), means_(means), var_(var),
        n_component_(n_component), sensor_location_(sensor_location) {
  }

  void setParam(std::vector<double> &&pi_vector,
                std::vector<Eigen::Vector3d> &&means,
                std::vector<Eigen::Matrix3d> &&var, int n_component,
                Eigen::Vector3d sensor_location) {
    pi_vector_ = pi_vector;
    means_ = means;
    var_ = var;
    n_component_ = n_component;
    sensor_location_ = sensor_location;

  }

  void setAgentId(int set_agent_id, int set_color_id) {
    agent_id = set_agent_id;
    color_id = set_color_id;
  }

  void GMM_fit_layer(std::vector<Eigen::Vector3d> &&point_cloud);

  std::vector<Eigen::Vector3d> GMMRandomSamplePoints(double n_sigma) const;

public:
  std::vector<double> pi_vector_;
  std::vector<Eigen::Vector3d> means_;
  std::vector<Eigen::Matrix3d> var_;
  int n_component_;
  Eigen::Vector3d sensor_location_;
  int agent_id{0};
  int color_id{37};
};

struct Ellipse {
  Ellipse(const Eigen::Vector3d &set_mean, const Eigen::Matrix3d &set_var)
      : mean(set_mean), var(set_var) {}
  Eigen::Vector3d mean;
  Eigen::Matrix3d var;
};

class EllipseFilter {
public:
  void SetParam(int SET_ELLIPISE_POOL_MAX_ELLIPSE_NUM,
                int SET_SIMILAR_NUM_THREDSHOLD,
                double SET_MEAN_SIMILAR_THRESHOLD,
                double SET_VAR_SIMILAR_THRESHOLD,
                double SET_FENCE_PLUS_X,
                double SET_FENCE_MINUS_X,
                double SET_FENCE_PLUS_Y,
                double SET_FENCE_MINUS_Y,
                double SET_FENCE_PLUS_Z,
                double SET_FENCE_MINUS_Z) {
    ELLIPISE_POOL_MAX_ELLIPSE_NUM = SET_ELLIPISE_POOL_MAX_ELLIPSE_NUM;
    SIMILAR_NUM_THREDSHOLD = SET_SIMILAR_NUM_THREDSHOLD;
    MEAN_SIMILAR_THRESHOLD = SET_MEAN_SIMILAR_THRESHOLD;
    VAR_SIMILAR_THRESHOLD = SET_VAR_SIMILAR_THRESHOLD;
    FENCE_PLUS_X=SET_FENCE_PLUS_X;
    FENCE_MINUS_X=SET_FENCE_MINUS_X;
    FENCE_PLUS_Y=SET_FENCE_PLUS_Y;
    FENCE_MINUS_Y=SET_FENCE_MINUS_Y;
    FENCE_PLUS_Z=SET_FENCE_PLUS_Z;
    FENCE_MINUS_Z=SET_FENCE_MINUS_Z;
  }
  bool FilterEllipse(Ellipse &ellipse) {
    if(ellipse.mean(0)>FENCE_PLUS_X || ellipse.mean(0)<FENCE_MINUS_X ||
       ellipse.mean(1)>FENCE_PLUS_Y || ellipse.mean(1)<FENCE_MINUS_Y ||
       ellipse.mean(2)>FENCE_PLUS_Z || ellipse.mean(2)<FENCE_MINUS_Z){
        return false;
    }
    for (auto it=merged_ellipse.begin(); it!=merged_ellipse.end(); it++) {
      if (isSimilar(ellipse, *it)) {
        Merge(ellipse, *it);
        return true;
      }
    }
    int similar_num = 0;
    for (auto v : ellipise_pool) {
      if (isSimilar(ellipse, v)) {
        similar_num++;
        if (similar_num > SIMILAR_NUM_THREDSHOLD) {
          merged_ellipse.push_back(ellipse);
          if(merged_ellipse.size()>100)
            merged_ellipse.pop_front();
          return true;
        }
      }
    }
    ellipise_pool.push_back(ellipse);
    if (ellipise_pool.size() > ELLIPISE_POOL_MAX_ELLIPSE_NUM)
      ellipise_pool.pop_front();
    return false;
  }

private:
  void Merge(Ellipse &ellipse1, Ellipse &ellipse2) {
    Eigen::Vector3d merged_mean = ellipse1.mean * 0.1 + ellipse2.mean * 0.9;
    Eigen::Matrix3d merged_var = ellipse1.var * 0.1 + ellipse2.var * 0.9;
    // ellipse1=Ellipse(merged_mean,merged_var);
    ellipse2 = Ellipse(merged_mean, merged_var);
  }

  bool isSimilar(const Ellipse &ellipse1, const Ellipse &ellipse2) {
    if ((ellipse1.mean - ellipse2.mean).norm() > MEAN_SIMILAR_THRESHOLD) {
      return false;
    }
    Eigen::EigenSolver<Eigen::Matrix3d> es1(ellipse1.var);
    Eigen::Matrix3d value_martrix1 = es1.pseudoEigenvalueMatrix();
    Eigen::EigenSolver<Eigen::Matrix3d> es2(ellipse2.var);
    Eigen::Matrix3d value_martrix2 = es2.pseudoEigenvalueMatrix();
    if (std::fabs(std::sqrt(value_martrix1(0, 0)) -
                  std::sqrt(value_martrix2(0, 0))) > VAR_SIMILAR_THRESHOLD ||
        std::fabs(std::sqrt(value_martrix1(1, 1)) -
                  std::sqrt(value_martrix2(1, 1))) > VAR_SIMILAR_THRESHOLD ||
        std::fabs(std::sqrt(value_martrix1(2, 2)) -
                  std::sqrt(value_martrix2(2, 2))) > VAR_SIMILAR_THRESHOLD) {
      return false;
    }
    return true;
  }
public:
  std::list<Ellipse> merged_ellipse;
  std::list<Ellipse> ellipise_pool;
  int ELLIPISE_POOL_MAX_ELLIPSE_NUM;
  int SIMILAR_NUM_THREDSHOLD;
  double MEAN_SIMILAR_THRESHOLD;
  double VAR_SIMILAR_THRESHOLD;
  double FENCE_PLUS_X;
  double FENCE_MINUS_X;
  double FENCE_PLUS_Y;
  double FENCE_MINUS_Y;
  double FENCE_PLUS_Z;
  double FENCE_MINUS_Z;
};

class GMMManager {
public:
  GMMManager() { tmp_pc.reset(new PC_TYPE()); }
  void SetParam(double SET_CLIP_DISTANCE, double SET_VOXEL_RES,
                double SET_EC_DIST, double SET_NOISE_RATIO_THRESHOLD,
                double SET_DETERMINANT_THRESHOLD,
                int SET_POINT_NUM_MINIMUM_THRESHOLD, double SET_UAV_BOX_X,
                double SET_UAV_BOX_Y, double SET_UAV_BOX_Z,
                int SET_ELLIPISE_POOL_MAX_ELLIPSE_NUM,
                int SET_SIMILAR_NUM_THREDSHOLD,
                double SET_MEAN_SIMILAR_THRESHOLD,
                double SET_VAR_SIMILAR_THRESHOLD,
                double SET_FENCE_PLUS_X,
                double SET_FENCE_MINUS_X,
                double SET_FENCE_PLUS_Y,
                double SET_FENCE_MINUS_Y,
                double SET_FENCE_PLUS_Z,
                double SET_FENCE_MINUS_Z) {
    CLIP_DISTANCE = SET_CLIP_DISTANCE;
    VOXEL_RES = SET_VOXEL_RES;
    EC_DIST = SET_EC_DIST;
    NOISE_RATIO_THRESHOLD = SET_NOISE_RATIO_THRESHOLD;
    DETERMINANT_THRESHOLD = SET_DETERMINANT_THRESHOLD;
    POINT_NUM_MINIMUM_THRESHOLD = SET_POINT_NUM_MINIMUM_THRESHOLD;
    UAV_BOX_X = SET_UAV_BOX_X;
    UAV_BOX_Y = SET_UAV_BOX_Y;
    UAV_BOX_Z = SET_UAV_BOX_Z;

    ellipse_filter.SetParam(
        SET_ELLIPISE_POOL_MAX_ELLIPSE_NUM, SET_SIMILAR_NUM_THREDSHOLD,
        SET_MEAN_SIMILAR_THRESHOLD, SET_VAR_SIMILAR_THRESHOLD, 
        SET_FENCE_PLUS_X, SET_FENCE_MINUS_X, SET_FENCE_PLUS_Y,
        SET_FENCE_MINUS_Y,SET_FENCE_PLUS_Z,SET_FENCE_MINUS_Z);
  }

  void TransferPointCloud(PC_PTR_TYPE input_pc,
                          const Eigen::Matrix4d &transform,
                          const std::vector<Eigen::Vector3d> &current_p =
                              std::vector<Eigen::Vector3d>(0));

  bool FitGMM(PC_PTR_TYPE input_pc, const Eigen::Vector3d &current_p,
              gvm::GMM &total_gmm);

  bool NaiveFitGMM(PC_PTR_TYPE input_pc, const Eigen::Vector3d &current_p,
                   gvm::GMM &total_gmm);

  void FilterGMM(gvm::GMM &gmm);

public:
  PC_PTR_TYPE tmp_pc;
  EllipseFilter ellipse_filter;
  double CLIP_DISTANCE;
  double VOXEL_RES;
  double EC_DIST;
  double NOISE_RATIO_THRESHOLD;
  double DETERMINANT_THRESHOLD;
  int POINT_NUM_MINIMUM_THRESHOLD;
  double UAV_BOX_X;
  double UAV_BOX_Y;
  double UAV_BOX_Z;
};

} // namespace gvm

#endif